package com.aliware.tianchi;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * Description:
 * Date: 19-7-1
 *
 * @author ujued
 */
public class WeightAdjust<K> {
    private Map<K, Integer> weightMap = new ConcurrentHashMap<>();
    private AtomicLong fullWeight = new AtomicLong(0);
    private Map<K, Integer> weightValMap = new ConcurrentHashMap<>();

    private Map<K, Integer> weightMapTmp = new HashMap<>();
    private AtomicLong fullWeightTmp = new AtomicLong(0);

    public void init(Map<K, Integer> init) {
        weightMapTmp.clear();
        fullWeightTmp.set(0);
        weightValMap = init;
        List<Entry<K, Integer>> l = init.entrySet().stream().sorted((e1, e2) -> e1.getValue() <= e2.getValue() ? -1 : 1).collect(Collectors.toList());
        for (int i = 0; i < l.size(); i++) {
            Entry<K, Integer> e = l.get(i);
            fullWeightTmp.set(fullWeightTmp.get() + e.getValue());
            if (i == 0) {
                weightMapTmp.put(e.getKey(), e.getValue());
                continue;
            }
            weightMapTmp.put(e.getKey(), e.getValue() + weightMapTmp.get(l.get(i - 1).getKey()));
        }
        weightMap.putAll(weightMapTmp);
        fullWeight.set(fullWeightTmp.get());
    }

    public K selectByWeight() {
        if (fullWeight.compareAndSet(0, 0)) {
            return null;
        }

        int selector = (int) (Math.random() * fullWeight.get());
        for (Entry<K, Integer> e : weightMap.entrySet().stream().sorted((e1, e2) -> e1.getValue() <= e2.getValue() ? -1 : 1).collect(Collectors.toList())) {
            if (e.getValue() > selector) {
                return e.getKey();
            }
        }
        throw new IllegalStateException("sys err.");
    }

    @Override
    public String toString() {
        return weightMap.toString();
    }

    public void adjustWeight(K invoker, Integer weight) {
        weightValMap.put(invoker, weight);
        init(weightValMap);
    }
}
